from langchain_core.tools import tool
import pandas as pd
import os
from datetime import datetime, timedelta
from django.db import connection
import json

@tool
def data_collection_tool(query: str) -> str:
    """
    收集司机历史接单数据并写入CSV文件。
    
    Args:
        query: JSON字符串，格式：{"days": 7, "output_file": "driver_data.csv", "time_category": "morning_peak"}
    
    Returns:
        收集的数据统计信息
    """
    try:
        params = json.loads(query)
        days = params.get('days', 7)
        output_file = params.get('output_file', 'driver_order_data.csv')
        time_category = params.get('time_category', None)

        # 收集数据：司机ID, 区域ID, 接单数量, 时段分类
        end_date = datetime.now()
        start_date = end_date - timedelta(days=days)

        # SQL查询历史订单数据 - 正确识别司机
        sql = """
        SELECT
            o.driver_id as driver_id,
            CONCAT(
                ROUND(COALESCE(o.origin_lng, 116.40), 2), '_',
                ROUND(COALESCE(o.origin_lat, 39.90), 2)
            ) as area_id,
            COUNT(*) as order_count,
            CASE
                WHEN HOUR(o.create_time) BETWEEN 7 AND 9 THEN 'morning_peak'
                WHEN HOUR(o.create_time) BETWEEN 17 AND 19 THEN 'evening_peak'
                WHEN HOUR(o.create_time) BETWEEN 21 AND 23 THEN 'night_peak'
                ELSE 'normal'
            END as time_category,
            ROUND(AVG(o.order_amount), 2) as avg_amount,
            COUNT(DISTINCT DATE(o.create_time)) as active_days
        FROM order_order o
        INNER JOIN users u ON o.driver_id = u.id
        WHERE o.create_time >= %s
            AND o.create_time <= %s
            AND o.driver_id IS NOT NULL
            AND u.type = 1
        """
        
        params_list = [start_date, end_date]
        
        # 如果指定了时段，先尝试查询该时段的数据
        if time_category and time_category != 'normal':
            # 先检查指定时段是否有数据（使用简化的查询）
            check_sql = """
            SELECT COUNT(*) as count
            FROM order_order o
            INNER JOIN users u ON o.driver_id = u.id
            WHERE o.create_time >= %s
                AND o.create_time <= %s
                AND o.driver_id IS NOT NULL
                AND u.type = 1
                AND CASE 
                    WHEN HOUR(o.create_time) BETWEEN 7 AND 9 THEN 'morning_peak'
                    WHEN HOUR(o.create_time) BETWEEN 17 AND 19 THEN 'evening_peak'
                    WHEN HOUR(o.create_time) BETWEEN 21 AND 23 THEN 'night_peak'
                    ELSE 'normal'
                END = %s
            """
            check_params = [start_date, end_date, time_category]
            
            with connection.cursor() as cursor:
                cursor.execute(check_sql, check_params)
                result = cursor.fetchone()
                count = result[0] if result else 0
            
            # 如果指定时段没有数据，回退到平时时段
            if count == 0:
                print(f"时段 {time_category} 没有数据，回退到平时时段")
                time_category = 'normal'
                params_list = [start_date, end_date]
            else:
                # 添加时段过滤条件到主查询
                sql += " AND CASE WHEN HOUR(o.create_time) BETWEEN 7 AND 9 THEN 'morning_peak' WHEN HOUR(o.create_time) BETWEEN 17 AND 19 THEN 'evening_peak' WHEN HOUR(o.create_time) BETWEEN 21 AND 23 THEN 'night_peak' ELSE 'normal' END = %s"
                params_list = [start_date, end_date, time_category]
        
        # 如果指定了平时时段或者回退到平时时段
        if time_category == 'normal':
            sql += " AND CASE WHEN HOUR(o.create_time) BETWEEN 7 AND 9 THEN 'morning_peak' WHEN HOUR(o.create_time) BETWEEN 17 AND 19 THEN 'evening_peak' WHEN HOUR(o.create_time) BETWEEN 21 AND 23 THEN 'night_peak' ELSE 'normal' END = 'normal'"
            
        sql += """
        GROUP BY o.driver_id, area_id, time_category
        HAVING order_count > 0
        ORDER BY o.driver_id, area_id, time_category
        """

        with connection.cursor() as cursor:
            cursor.execute(sql, params_list)
            results = cursor.fetchall()

        # 转换为DataFrame并保存CSV
        columns = ['driver_id', 'area_id', 'order_count', 'time_category', 'avg_amount', 'active_days']
        df = pd.DataFrame(results, columns=columns)

        # 检查是否有数据
        if df.empty:
            return f"数据收集失败：在指定时段 {time_category} 中没有找到订单数据"

        # 确保driver_id是整数类型
        df['driver_id'] = df['driver_id'].astype(int)

        # 确保数据目录存在
        os.makedirs('ai_data', exist_ok=True)
        file_path = f'ai_data/{output_file}'
        df.to_csv(file_path, index=False, encoding='utf-8')

        stats = {
            'total_records': len(df),
            'unique_drivers': df['driver_id'].nunique(),
            'unique_areas': df['area_id'].nunique(),
            'time_categories': df['time_category'].value_counts().to_dict(),
            'file_path': file_path,
            'collection_date': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            'query_params': {
                'days': days,
                'start_date': start_date.strftime('%Y-%m-%d'),
                'end_date': end_date.strftime('%Y-%m-%d'),
                'time_category': time_category,
                'actual_time_category': df['time_category'].iloc[0] if len(df) > 0 else 'normal'
            }
        }

        return f"数据收集完成：{json.dumps(stats, ensure_ascii=False, indent=2)}"

    except Exception as e:
        return f"数据收集失败：{str(e)}"
